package top.luoyuanxiang.personalknowledgebase.controller;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.reader.pdf.PagePdfDocumentReader;
import org.springframework.ai.reader.pdf.ParagraphPdfDocumentReader;
import org.springframework.ai.reader.tika.TikaDocumentReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.multipart.MultipartFile;
import top.luoyuanxiang.personalknowledgebase.service.IDocumentEmbeddingService;

import java.util.List;

/**
 * RAG增强
 *
 * @author luoyuanxiang
 */
@CrossOrigin("*")
@RestController
@RequestMapping("/documentEmbedding")
public class DocumentEmbeddingController {

    @Resource
    private IDocumentEmbeddingService documentEmbeddingService;

    /**
     * PagePdfDocumentReader读取PDF信息(spring-ai-pdf-document-reader)
     *
     * @param file 文件
     * @return {@link ResponseEntity }<{@link List }<{@link Document }>>
     */
    @PostMapping("/pagePdfDocumentReader")
    public ResponseEntity<List<Document>> pagePdfDocumentReader(MultipartFile file) {
        var pagePdfDocumentReader = new PagePdfDocumentReader(file.getResource());
        return ResponseEntity.ok(pagePdfDocumentReader.get());
    }

    /**
     * ParagraphPdfDocumentReader读取PDF信息(spring-ai-pdf-document-reader)
     *
     * @param file 文件
     * @return {@link ResponseEntity }<{@link List }<{@link Document }>>
     */
    @PostMapping("/paragraphPdfDocumentReader")
    public ResponseEntity<List<Document>> paragraphPdfDocumentReader(MultipartFile file) {
        var paragraphPdfDocumentReader = new ParagraphPdfDocumentReader(file.getResource());
        return ResponseEntity.ok(paragraphPdfDocumentReader.get());
    }

    /**
     * tikaDocumentReader读取文档信息(spring-ai-tika-document-reader)
     *
     * @param file 文件
     * @return {@link ResponseEntity }<{@link List }<{@link Document }>>
     */
    @PostMapping("/tikaDocumentReader")
    public ResponseEntity<Integer> tikaDocumentReader(MultipartFile file) {
        var tikaDocumentReader = new TikaDocumentReader(file.getResource());
        // 2. split trunks
        List<Document> splitDocuments = TokenTextSplitter.builder()
                .build()
                .apply(tikaDocumentReader.get());
        return ResponseEntity.ok(splitDocuments.size());
    }

    /**
     * 添加文档
     *
     * @param files 文件
     */
    @PostMapping("/addDocument")
    public ResponseEntity<Boolean> addDocument(MultipartFile[] files) {
        documentEmbeddingService.addDocument(files);
        return ResponseEntity.ok(true);
    }

    /**
     * 搜索文档
     *
     * @param query 查询
     * @return {@link List }<{@link Document }>
     */
    @GetMapping("/searchDocument")
    public ResponseEntity<List<Document>> searchDocument(String query){
        return ResponseEntity.ok(documentEmbeddingService.searchDocument(query));
    }
}
